#!/usr/bin/env python3
import argparse
import os
from typing import Dict, List, Optional, Set

import yaml


def load_yaml(file_path: str) -> Dict:
    with open(file_path, 'r') as f:
        return yaml.safe_load(f)


def validate_business_codes(metadata: Dict) -> None:
    """Validate that all business codes are unique within each app."""
    for app in metadata['app']:
        seen_codes: Set[int] = set()
        for biz in app['business']:
            code = biz['code']
            if code in seen_codes:
                raise ValueError(
                    f"Duplicate business code {code} found for business '{biz['name']}' in app '{app['name']}'")
            seen_codes.add(code)


def get_biz_code(metadata: Dict, biz_name: str) -> Optional[int]:
    for app in metadata['app']:
        if app['name'] == 'cozeloop':
            for biz in app['business']:
                if biz['name'] == biz_name:
                    return biz['code']
    return None


def error_code(app_code: int, biz_code: int, sub_code: int) -> int:
    return app_code * 100000000 + biz_code * 10000 + sub_code


def generate_go_code(biz_name: str, biz_code: int, common_errors: List[Dict], biz_errors: List[Dict],
                     output_dir: Optional[str] = None) -> str:
    if not output_dir:
        raise ValueError("output_dir cannot be empty")

    app_code = 6  # cozeloop
    package_name = biz_name.lower()

    # Generate constants
    constants = []
    registrations = []

    # Add common errors
    # Build up all constants first
    const_lines = []
    for error in common_errors:
        code = error_code(app_code, biz_code, error['code'])
        name = error['name']
        unexport_name = name[0].lower() + name[1:]
        message = error.get('message', '')
        description = error.get('description', '')
        no_affect = error.get('no_affect_stability', False)

        const_lines.extend([
            f"\t{name}Code = {code}{' // ' + description if description else ''}",
            f"\t{unexport_name}Message = \"{message}\"",
            f"\t{unexport_name}NoAffectStability = {str(no_affect).lower()}",
            ""  # Add empty line between errors
        ])
        registrations.append(
            f"\tcode.Register(\n\t\t{name}Code,\n\t\t{unexport_name}Message,\n\t\tcode.WithAffectStability(!{unexport_name}NoAffectStability),\n\t)\n")

    # Add business specific errors
    if biz_errors:
        for error in biz_errors:
            code = error_code(app_code, biz_code, error['code'])
            name = error['name']
            unexport_name = name[0].lower() + name[1:]
            message = error.get('message', '')
            description = error.get('description', '')
            no_affect = error.get('no_affect_stability', False)

            const_lines.extend([
                f"\t{name}Code = {code}{' // ' + description if description else ''}",
                f"\t{unexport_name}Message = \"{message}\"",
                f"\t{unexport_name}NoAffectStability = {str(no_affect).lower()}",
                ""  # Add empty line between errors
            ])
            registrations.append(
                f"\tcode.Register(\n\t\t{name}Code,\n\t\t{unexport_name}Message,\n\t\tcode.WithAffectStability(!{unexport_name}NoAffectStability),\n\t)\n")

    # Remove the last empty line if it exists
    if const_lines and const_lines[-1] == "":
        const_lines.pop()

    # Combine all constants into a single const block
    constants = [f"const (\n{chr(10).join(const_lines)}\n)"]

    output_path = os.path.join(output_dir, f"{package_name}.go")

    # Generate the complete Go file
    go_code = (
            f"""// Code generated by tool. DO NOT EDIT.
// app: cozeloop, biz: {biz_name}

package {os.path.basename(os.path.dirname(output_path))}

import (
""" + '\t"github.com/coze-dev/coze-loop/backend/pkg/errorx/code"\n' + """
)

""" + '\n'.join(constants) + """

func init() {
""" + ''.join(['\n' + r for r in registrations]) + """
}
"""
    )

    os.makedirs(output_dir, exist_ok=True)
    with open(output_path, 'w') as f:
        f.write(go_code)
    # Run go fmt after generating the file
    # Run go fmt and check return code
    ret = os.system(f"go fmt {output_path}")
    if ret != 0:
        print(f"Warning: go fmt failed with return code {ret}")
    return output_path


def generate_biz_code(biz_name: str, biz_code: int, common_errors: List[Dict], output_dir: Optional[str] = None) -> str:
    # Get business specific errors if they exist
    biz_errors = []
    biz_error_file = os.path.join(os.path.dirname(__file__), f'{biz_name}.yaml')
    if os.path.exists(biz_error_file):
        biz_errors = load_yaml(biz_error_file)['error_code']

    # Generate and output code
    current_file_dir = os.path.dirname(os.path.abspath(__file__))
    project_dir = os.path.dirname(os.path.dirname(os.path.dirname(current_file_dir)))
    if not output_dir:
        output_dir = os.path.join(project_dir, 'backend/modules', biz_name, 'pkg/errno')
    else:
        output_dir = os.path.expandvars(output_dir)
        if not os.path.isabs(output_dir):
            output_dir = os.path.join(project_dir, output_dir)

    return generate_go_code(biz_name, biz_code, common_errors, biz_errors, output_dir)


def main():
    parser = argparse.ArgumentParser(description='Generate error codes for a specific business domain')
    parser.add_argument('biz', help='Business domain name (e.g., evaluation) or "*" to generate for all business domains')
    parser.add_argument('--output-dir', help='Output directory for generated Go file')
    args = parser.parse_args()

    # Load configuration files
    metadata = load_yaml(os.path.join(os.path.dirname(__file__), 'metadata.yaml'))

    # Validate business codes before proceeding
    try:
        validate_business_codes(metadata)
    except ValueError as e:
        print(f"Error in metadata.yaml: {e}")
        return

    common_errors = load_yaml(os.path.join(os.path.dirname(__file__), 'common.yaml'))['error_code']

    # Get cozeloop app from metadata
    cozeloop_app = None
    for app in metadata['app']:
        if app['name'] == 'cozeloop':
            cozeloop_app = app
            break

    if not cozeloop_app:
        print("Error: cozeloop app not found in metadata")
        return

    # Handle wildcard case
    if args.biz == '*':
        for biz in cozeloop_app['business']:
            biz_name = biz['name']
            if biz_name == 'common':
                continue

            print(f"\nProcessing business domain: {biz_name}")
            result = generate_biz_code(biz_name, biz['code'], common_errors, args.output_dir)
            print(f"Generated error codes written to: {result}")
        return

    # Handle single business domain case
    biz_code = get_biz_code(metadata, args.biz)
    if biz_code is None:
        print(f"Error: Business domain '{args.biz}' not found in metadata")
        return

    result = generate_biz_code(args.biz, biz_code, common_errors, args.output_dir)
    print(f"Generated error codes written to: {result}")


if __name__ == '__main__':
    main()
